import datetime
import logging
import os
from collections import defaultdict
from dataclasses import dataclass
from typing import Dict, List, Optional, Set, Tuple

import pytz
from absl import app, flags

from xunkemgmt_client.base.datetime_util import to_timestamp_int
from xunkemgmt_client.support.pta.plot.plot_util import StrategyStatLogParser

import matplotlib

matplotlib.use('Agg')
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import pandas as pd

from coin.proto.coin_strategy_summary_pb2 import StrategyStatLog

FLAGS = flags.FLAGS


@dataclass
class StatLogParam:
  strat_name: str
  accounting_currency: str
  start_time: datetime.datetime
  end_time: datetime.datetime
  plot_dir: str
  strategy_rows: List[List[Optional[float]]] = None
  symbol_map: Dict[Tuple, List[List[Optional[float]]]] = None


def _parse_plot_df(rows: List[List], field_name: str) -> Tuple[List, List]:
  valid_ts_rows = []
  valid_data_rows = []
  idx_map = {
      "strategy_pnl": 1,
      "strategy_pnl_from_fill": 2,
      "symbol_pnl": 1,
      "symbol_position": 2,
      "symbol_position_from_fill": 3,
  }
  idx = idx_map[field_name]
  df = pd.DataFrame(rows)
  if len(df) == 0:
    return valid_ts_rows, valid_data_rows
  ts_rows = [datetime.datetime.utcfromtimestamp(ts / 1e+9) for ts in df.iloc[:, 0].tolist()]
  data_rows = df.iloc[:, idx].tolist()
  for idx in range(len(ts_rows)):
    if pd.isna(data_rows[idx]):
      continue
    valid_ts_rows.append(ts_rows[idx])
    valid_data_rows.append(data_rows[idx])
  return valid_ts_rows, valid_data_rows


class StrategyPlotter(object):
  def __init__(self,
               account_df: pd.DataFrame,
               pnl_adj_info: Optional[Dict[str, List[Dict]]]=None):
    self.time_fmt = "%Y%m%dT%H%M%S"
    self._account_df = account_df
    self._pnl_adj_info = pnl_adj_info
    self._pnl_adjusted_date: Set[str] = set()

  def _setup_plot(self, figsize_y=None):
    plt.style.use('seaborn-v0_8-whitegrid')
    plt.rcParams['image.cmap'] = 'tab20c'
    plt.rcParams['lines.linewidth'] = 0.5
    plt.rcParams['figure.dpi'] = 200
    plt.rcParams['figure.figsize'] = [8, figsize_y or 3]
    plt.rcParams['font.family'] = ['sans-serif']
    plt.rcParams['font.size'] = 6
    plt.rcParams['legend.fontsize'] = 6
    plt.rcParams['xtick.labelsize'] = 6
    plt.rcParams['ytick.labelsize'] = 6
    plt.rcParams['savefig.facecolor'] = '#ffffff'

  def _get_account_meo(self, account_id: int):
    df = self._account_df[self._account_df["account_id"] == account_id].iloc[0]
    return f"{df.exchange}_{df.owner}"

  def _patch_symbol_keys(self, symbol_keys: List[Tuple]) -> List[Tuple]:
    result = []
    for symbol_key in symbol_keys:
      symbol, account_id = symbol_key
      account_meo = self._get_account_meo(account_id)
      result.append((symbol, account_id, account_meo))
    return result

  def _do_pnl_plot(
      self,
      stat_log_param: StatLogParam,
  ):
    strat_name = stat_log_param.strat_name
    accounting_currency = stat_log_param.accounting_currency
    start_time = stat_log_param.start_time
    end_time = stat_log_param.end_time
    symbol_keys = list(stat_log_param.symbol_map.keys())
    symbol_keys = self._patch_symbol_keys(symbol_keys)
    plot_dir = stat_log_param.plot_dir

    subplot_rows = 2 + len(symbol_keys)
    fig_height = subplot_rows if subplot_rows > 6.5 else 6.5
    self._setup_plot(fig_height)
    fig = plt.figure()
    ax1 = plt.subplot2grid((subplot_rows, 1), (0, 0), rowspan=2)
    times_1, values_1 = _parse_plot_df(stat_log_param.strategy_rows, "strategy_pnl")
    times_2, values_2 = _parse_plot_df(stat_log_param.strategy_rows, "strategy_pnl_from_fill")
    ax1.plot(times_1, values_1, label=strat_name)
    ax1.plot(times_2, values_2, label=f'{strat_name} (from_fill)', linestyle=':')
    label_list = [strat_name]
    symbol_keys = sorted(symbol_keys, key=lambda x: (x[0], x[1]))
    for i, symbol_key in enumerate(symbol_keys):
      symbol, account_id, account_meo = symbol_key
      symbol_rows = stat_log_param.symbol_map[(symbol, account_id)]
      symbol_label = f"{symbol}_{account_meo}"
      times, values = _parse_plot_df(symbol_rows, "symbol_pnl")
      ax1.plot(times, values, label=symbol_label)
      label_list.append(symbol_label)
      times_1, values_1 = _parse_plot_df(symbol_rows, "symbol_position")
      times_2, values_2 = _parse_plot_df(symbol_rows, "symbol_position_from_fill")
      ax2 = plt.subplot2grid((subplot_rows, 1), (2 + i, 0), sharex=ax1)
      ax2.plot(times_1, values_1, label=symbol_label, drawstyle='steps-post')
      ax2.plot(times_2, values_2, label=symbol_label, linestyle=':', drawstyle='steps-post')
      ax2.legend([symbol_label], loc='upper left')

    ax1.set_xlim((start_time, end_time))
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d %HH'))
    fig.autofmt_xdate()
    ax1.legend(loc='lower left', bbox_to_anchor=(0.0, 1.01), borderaxespad=0, ncol=3)
    fig.tight_layout(pad=0.5)
    filename = '-'.join([
        strat_name,
        accounting_currency,
    ])
    filepath = os.path.join(plot_dir, f"{filename}.png")
    fig.savefig(filepath, transparent=False)
    plt.close(fig)

  def _check_empty_rows(self, rows: List[List]) -> bool:
    for row in rows:
      if len(row) == 0:
        continue
      for elem in row[1:]:
        if elem is not None and elem != 0.0:
          return False
    return True

  def _get_pnl_adjust_value(
      self, strat_name: str, ccy: str, ts: int, curr_pnl: float, prev_pnl: float):
    ret = None
    if self._pnl_adj_info is not None:
      td_str = datetime.datetime.utcfromtimestamp(ts / 10**9).date().strftime('%Y%m%d')
      if td_str not in self._pnl_adjusted_date:
        adj_entries = self._pnl_adj_info.get(td_str, [])
        adj_values = [
            entry['pnl_adj'] for entry in adj_entries
            if strat_name == entry['strategy_name'] and
               ccy == entry['accounting_currency']
        ]
        if len(adj_values) > 0:
          assert len(adj_values) == 1, len(adj_values)
          adj_value = adj_values[0]
          if 0.95 < (prev_pnl - curr_pnl) / adj_value < 1.05:
            ret = adj_value
            self._pnl_adjusted_date.add(td_str)
            logging.info(
                'td: %s, strategy: %s, currency: %s, pnl_adj: %s, ts: %s, curr_pnl: %s, prev_pnl: %s'
                % (td_str,
                   strat_name,
                   ccy,
                   adj_value,
                   ts,
                   curr_pnl,
                   prev_pnl))
    return ret

  def _get_adjusted_rows(self, strat_name: str, ccy: str, rows: List[List]):
    total_adj = 0
    prev_pnl = 0
    for row in rows:
      curr_pnl = row[1]
      if curr_pnl is None:
        continue
      adj = self._get_pnl_adjust_value(
          strat_name=strat_name, ccy=ccy, ts=row[0], curr_pnl=row[1], prev_pnl=prev_pnl)
      prev_pnl = curr_pnl
      if adj:
        total_adj += adj
      row[1] += total_adj
    return rows

  def do_pnl_plot(
      self,
      strat_stat_logs: List[StrategyStatLog],
      strat_name: str,
      start_time: datetime.datetime,
      end_time: datetime.datetime,
      plot_dir: str,
  ):
    start_ts = to_timestamp_int(start_time)
    end_ts = to_timestamp_int(end_time)
    os.makedirs(plot_dir, exist_ok=True)

    strat_stat_log_parser = StrategyStatLogParser(start_ts=start_ts,)
    for pb in strat_stat_logs:
      if pb.strategy_stat.stat_ts < start_ts \
        or pb.strategy_stat.stat_ts > end_ts:
        continue
      strat_stat_log_parser.process_pb(pb)
    strategy_data = strat_stat_log_parser.strategy_data
    symbol_data = strat_stat_log_parser.symbol_data

    currency_map = {}
    for ccy, rows in strategy_data.items():
      if self._check_empty_rows(rows):
        logging.info('ignore empty strat data: %s, %s' % (strat_name, ccy))
        continue
      adjusted_rows = self._get_adjusted_rows(strat_name, ccy, rows)
      if ccy not in currency_map:
        currency_map[ccy] = StatLogParam(
            strat_name=strat_name,
            accounting_currency=ccy,
            symbol_map=defaultdict(list),
            start_time=start_time,
            end_time=end_time,
            plot_dir=plot_dir,
        )
      currency_map[ccy].strategy_rows = adjusted_rows
    for ccy, value in symbol_data.items():
      for symbol_key, rows in value.items():
        if self._check_empty_rows(rows):
          logging.info('ignore empty symbol data: %s, %s' % (strat_name, symbol_key))
          continue
        if ccy not in currency_map:
          currency_map[ccy] = StatLogParam(
              strat_name=strat_name,
              accounting_currency=ccy,
              symbol_map=defaultdict(list),
              start_time=start_time,
              end_time=end_time,
              plot_dir=plot_dir,
          )
        currency_map[ccy].symbol_map[symbol_key] = rows

    for _, stat_log_param in currency_map.items():
      self._do_pnl_plot(stat_log_param)


def main(_):
  from coin.support.pta.logging.strategy_stat_logger import StrategyStatReader
  from coin.tool.strat_monitor.util.monitor_util import (
      get_pnl_adjust_info)
  from xunkemgmt_client.client.util.query_util import query_accounts

  class Onlog:
    def __init__(self):
      self.strat_stat_logs = []

    def on_log(self, timestamp: int, log):
      pb = StrategyStatLog()
      pb.ParseFromString(log)
      self.strat_stat_logs.append(pb)

  time_fmt = "%Y%m%dT%H%M%S"
  strategy_name = FLAGS.strategy_name
  machine = FLAGS.machine
  start_time = datetime.datetime.strptime(
      FLAGS.start_time, time_fmt).replace(tzinfo=pytz.UTC)
  end_time = datetime.datetime.strptime(
      FLAGS.end_time, time_fmt).replace(tzinfo=pytz.UTC)
  pnl_adj = {}
  if FLAGS.pnl_adj_file is not None:
    pnl_adj = get_pnl_adjust_info(FLAGS.pnl_adj_file, start_time.date(), end_time.date())
  root_dir = FLAGS.root_dir
  plot_dir = FLAGS.plot_dir

  onlog = Onlog()
  stat_reader = StrategyStatReader()
  stat_reader.run_from_archive(
      on_log_callback=onlog.on_log,
      start_time=start_time,
      end_time=end_time,
      root_dir=root_dir,
      machine=machine,
      strategy_name=strategy_name,
  )

  account_df = query_accounts()
  plotter = StrategyPlotter(account_df, pnl_adj)
  plotter.do_pnl_plot(
      strat_name=strategy_name,
      strat_stat_logs=onlog.strat_stat_logs,
      start_time=start_time,
      end_time=end_time,
      plot_dir=plot_dir,
  )


if __name__ == '__main__':
  """
./pyrunner python/coin/support/pta/plot/strat_plotter.py \
  --root_dir /remote/iosg/strat-1/buckets/analysis.derived.coin/live/pnl_stat_proto_log \
  --plot_dir /tmp \
  --machine strategy-26.ap-northeast-1.megaport \
  --strategy_name basis_smm2_binance_perpetual_binance_busd_perpetual \
  --start_time 20230307T000000 \
  --end_time 20230309T000000
  """

  flags.DEFINE_string("strategy_name", None, "strategy_name")
  flags.DEFINE_string("machine", None, "machine")
  flags.DEFINE_string("start_time", None, "%Y%m%dT%H%M%S")
  flags.DEFINE_string('pnl_adj_file', None, 'pnl_adj_file')
  flags.DEFINE_string("end_time", None, "%Y%m%dT%H%M%S")
  flags.DEFINE_string("root_dir", None, "root_dir")
  flags.DEFINE_string("plot_dir", None, "plot_dir")

  app.run(main)
